import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import random
from collections import deque

class PCBLayout:
    def __init__(self, grid_width=20, grid_height=20, grid_size=1.0):
        """
        初始化PCB布局网格
        :param grid_width: 网格宽度（单位网格数）
        :param grid_height: 网格高度（单位网格数）
        :param grid_size: 每个网格的实际大小（毫米）
        """
        self.grid_width = grid_width
        self.grid_height = grid_height
        self.grid_size = grid_size
        self.grid = np.zeros((grid_height, grid_width), dtype=int)
        self.components = []
        self.connections = []
        
        # 方向：上、右、下、左
        self.directions = [(-1, 0), (0, 1), (1, 0), (0, -1)]
    
    def add_component(self, component_id, width, height, pins=[]):
        """
        添加一个元件到布局中
        :param component_id: 元件唯一标识
        :param width: 元件宽度（网格数）
        :param height: 元件高度（网格数）
        :param pins: 引脚列表 [(相对x, 相对y)]
        """
        self.components.append({
            'id': component_id,
            'width': width,
            'height': height,
            'pins': pins,
            'position': None  # 将在布局时设置
        })
    
    def add_connection(self, component_id1, pin_id1, component_id2, pin_id2):
        """
        添加两个引脚之间的连接
        :param component_id1: 第一个元件ID
        :param pin_id1: 第一个元件的引脚索引
        :param component_id2: 第二个元件ID
        :param pin_id2: 第二个元件的引脚索引
        """
        self.connections.append({
            'comp1': component_id1,
            'pin1': pin_id1,
            'comp2': component_id2,
            'pin2': pin_id2
        })
    
    def is_valid_position(self, x, y, width, height):
        """检查位置是否有效（不超出边界且不与其他元件重叠）"""
        # 检查边界
        if x < 0 or y < 0 or x + width > self.grid_width or y + height > self.grid_height:
            return False
        
        # 检查重叠
        for i in range(height):
            for j in range(width):
                if self.grid[y + i, x + j] != 0:
                    return False
        return True
    
    def place_component(self, component, x, y):
        """将元件放置在指定位置"""
        width = component['width']
        height = component['height']
        
        if not self.is_valid_position(x, y, width, height):
            return False
        
        # 标记网格
        for i in range(height):
            for j in range(width):
                self.grid[y + i, x + j] = 1  # 简单标记为占用
        
        # 更新元件位置
        component['position'] = (x, y)
        return True
    
    def find_empty_position(self, width, height):
        """在网格中查找适合指定尺寸的空位置（从左到右，从上到下扫描）"""
        for y in range(self.grid_height - height + 1):
            for x in range(self.grid_width - width + 1):
                if self.is_valid_position(x, y, width, height):
                    return x, y
        return None
    
    def calculate_connection_length(self, connection):
        """计算两个引脚之间的曼哈顿距离"""
        comp1 = next(c for c in self.components if c['id'] == connection['comp1'])
        comp2 = next(c for c in self.components if c['id'] == connection['comp2'])
        
        if comp1['position'] is None or comp2['position'] is None:
            return float('inf')
        
        x1, y1 = comp1['position']
        pin1_x, pin1_y = comp1['pins'][connection['pin1']]
        abs_x1 = x1 + pin1_x
        abs_y1 = y1 + pin1_y
        
        x2, y2 = comp2['position']
        pin2_x, pin2_y = comp2['pins'][connection['pin2']]
        abs_x2 = x2 + pin2_x
        abs_y2 = y2 + pin2_y
        
        return abs(abs_x1 - abs_x2) + abs(abs_y1 - abs_y2)
    
    def calculate_total_connection_length(self):
        """计算所有连接的总长度"""
        total = 0
        for conn in self.connections:
            total += self.calculate_connection_length(conn)
        return total
    
    def optimize_layout(self, max_iterations=100):
        """优化布局以减少连接长度"""
        current_length = self.calculate_total_connection_length()
        
        for _ in range(max_iterations):
            # 随机选择两个元件尝试交换位置
            comp1, comp2 = random.sample([c for c in self.components if c['position']], 2)
            
            # 保存原始位置
            pos1 = comp1['position']
            pos2 = comp2['position']
            
            # 临时移除元件
            self.remove_component_from_grid(comp1)
            self.remove_component_from_grid(comp2)
            
            # 尝试交换位置
            success = False
            if self.place_component(comp1, pos2[0], pos2[1]):
                if self.place_component(comp2, pos1[0], pos1[1]):
                    success = True
            
            # 如果交换失败或没有改善，恢复原位置
            new_length = self.calculate_total_connection_length()
            if not success or new_length > current_length:
                # 恢复原始位置
                self.remove_component_from_grid(comp1)
                self.remove_component_from_grid(comp2)
                self.place_component(comp1, pos1[0], pos1[1])
                self.place_component(comp2, pos2[0], pos2[1])
            else:
                current_length = new_length
                self.visualize()
    
    def remove_component_from_grid(self, component):
        """从网格中移除元件"""
        if component['position'] is None:
            return
        
        x, y = component['position']
        width = component['width']
        height = component['height']
        
        for i in range(height):
            for j in range(width):
                self.grid[y + i, x + j] = 0
        
        component['position'] = None
    
    def auto_layout(self):
        """自动布局元件"""
        # 按元件大小排序（从大到小）
        sorted_components = sorted(self.components, key=lambda c: c['width'] * c['height'], reverse=True)
        
        # 放置元件
        for comp in sorted_components:
            pos = self.find_empty_position(comp['width'], comp['height'])
            if pos:
                self.place_component(comp, pos[0], pos[1])
            else:
                print(f"警告: 无法放置元件 {comp['id']}")
            self.visualize()
        
        # 优化布局
        self.optimize_layout()
    
    def visualize(self):
        """可视化PCB布局"""
        fig, ax = plt.subplots(figsize=(12, 12))
        
        # 绘制网格
        for x in range(self.grid_width + 1):
            ax.axvline(x * self.grid_size, color='gray', linestyle='-', alpha=0.3)
        for y in range(self.grid_height + 1):
            ax.axhline(y * self.grid_size, color='gray', linestyle='-', alpha=0.3)
        
        # 绘制元件
        colors = plt.cm.tab10.colors
        for i, comp in enumerate(self.components):
            if comp['position'] is None:
                continue
                
            x, y = comp['position']
            width = comp['width']
            height = comp['height']
            
            # 绘制元件边框
            rect = patches.Rectangle(
                (x * self.grid_size, y * self.grid_size),
                width * self.grid_size,
                height * self.grid_size,
                linewidth=2,
                edgecolor=colors[i % len(colors)],
                facecolor=colors[i % len(colors)] + (0.2,)
            )
            ax.add_patch(rect)
            
            # 绘制元件ID
            ax.text(
                (x + width/2) * self.grid_size,
                (y + height/2) * self.grid_size,
                comp['id'],
                ha='center',
                va='center',
                fontsize=10,
                fontweight='bold'
            )
            
            # 绘制引脚
            for pin_x, pin_y in comp['pins']:
                ax.plot(
                    (x + pin_x) * self.grid_size,
                    (y + pin_y) * self.grid_size,
                    'ko',
                    markersize=4
                )
        
        # 绘制连接
        for conn in self.connections:
            comp1 = next(c for c in self.components if c['id'] == conn['comp1'])
            comp2 = next(c for c in self.components if c['id'] == conn['comp2'])
            
            if comp1['position'] is None or comp2['position'] is None:
                continue
                
            x1, y1 = comp1['position']
            pin1_x, pin1_y = comp1['pins'][conn['pin1']]
            abs_x1 = (x1 + pin1_x) * self.grid_size
            abs_y1 = (y1 + pin1_y) * self.grid_size
            
            x2, y2 = comp2['position']
            pin2_x, pin2_y = comp2['pins'][conn['pin2']]
            abs_x2 = (x2 + pin2_x) * self.grid_size
            abs_y2 = (y2 + pin2_y) * self.grid_size
            
            ax.plot(
                [abs_x1, abs_x2],
                [abs_y1, abs_y2],
                'b-',
                linewidth=1,
                alpha=0.5
            )
        
        # 设置图形属性
        ax.set_xlim(0, self.grid_width * self.grid_size)
        ax.set_ylim(0, self.grid_height * self.grid_size)
        ax.set_aspect('equal')
        ax.set_title('PCB 自动布局')
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        ax.grid(True, linestyle='-', alpha=0.3)
        
        plt.tight_layout()
        plt.show()


# 示例使用
if __name__ == "__main__":
    # 创建PCB布局网格 (20x20 网格，每个网格1mm)
    pcb = PCBLayout(grid_width=20, grid_height=20, grid_size=1.0)
    
    # 添加元件（ID, 宽度, 高度, 引脚位置列表）
    pcb.add_component("CPU", 4, 4, [(1, 1), (1, 3), (3, 1), (3, 3)])
    pcb.add_component("RAM", 2, 4, [(0, 1), (0, 3), (2, 1), (2, 3)])
    pcb.add_component("GPU", 4, 3, [(1, 0), (1, 2), (3, 0), (3, 2)])
    pcb.add_component("SSD", 3, 2, [(0, 0), (0, 1), (2, 0), (2, 1)])
    pcb.add_component("NIC", 3, 2, [(0, 0), (0, 1), (2, 0), (2, 1)])
    pcb.add_component("USB1", 2, 1, [(0, 0), (1, 0)])
    pcb.add_component("USB2", 2, 1, [(0, 0), (1, 0)])
    pcb.add_component("AUDIO", 2, 2, [(0, 0), (0, 1), (1, 0), (1, 1)])
    pcb.add_component("PWR", 3, 2, [(1, 0), (1, 1)])
    
    # 添加连接（元件1ID, 引脚索引, 元件2ID, 引脚索引）
    pcb.add_connection("CPU", 0, "RAM", 0)
    pcb.add_connection("CPU", 1, "RAM", 1)
    pcb.add_connection("CPU", 2, "GPU", 0)
    pcb.add_connection("CPU", 3, "GPU", 1)
    pcb.add_connection("CPU", 0, "SSD", 0)
    pcb.add_connection("SSD", 1, "NIC", 0)
    pcb.add_connection("NIC", 1, "USB1", 0)
    pcb.add_connection("NIC", 2, "USB2", 0)
    pcb.add_connection("PWR", 0, "CPU", 0)
    pcb.add_connection("PWR", 1, "GPU", 2)
    pcb.add_connection("AUDIO", 0, "USB1", 1)
    pcb.add_connection("AUDIO", 1, "USB2", 1)
    
    # 执行自动布局
    pcb.auto_layout()
    
    # 计算总连接长度
    total_length = pcb.calculate_total_connection_length()
    print(f"总连接长度: {total_length} 网格单位")
    
    # 可视化布局
    pcb.visualize()